import dlib
import numpy as np
import cv2
import os
import psycopg2

def extract_face_features(dataset_path):
    proto_path = 'path_to_proto_file'
    model_path = 'path_to_model_file'

    predictor_path = "path_to_predictor_file"
    face_model_path = "path_to_face_model_file"

    face_net = cv2.dnn.readNetFromCaffe(proto_path, model_path)

    shape_predictor = dlib.shape_predictor(predictor_path)
    face_model = dlib.face_recognition_model_v1(face_model_path)

    image_folder = dataset_path

    db_conn = psycopg2.connect(
        host="YOUR_HOST",
        port="YOUR_PORT",
        user="YOUR_USER",
        password="YOUR_PASSWORD",
        database="YOUR_DATABASE"
    )

    db_cursor = db_conn.cursor()

    drop_table_query = """
        DROP TABLE IF EXISTS feature_vectors;
        """
    db_cursor.execute(drop_table_query)
    db_conn.commit()

    # Create the table
    create_table_query = """
    CREATE TABLE feature_vectors (
        id SERIAL PRIMARY KEY,
        filename VARCHAR(255),
        vector TEXT,
        distance REAL
    );
    """
    db_cursor.execute(create_table_query)
    db_conn.commit()

    for filename in os.listdir(image_folder):
        image_path = os.path.join(image_folder, filename)
        image = cv2.imread(image_path)
        if image is None:
            print(f"Unable to read image file: {image_path}")
            continue

        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        (height, width) = image.shape[:2]

        blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0, (300, 300), (104.0, 177.0, 123.0))
        face_net.setInput(blob)
        detections = face_net.forward()

        x_start, y_start, x_end, y_end = 0, 0, 0, 0
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([width, height, width, height])
                (x_start, y_start, x_end, y_end) = box.astype("int")
                break

        rect = dlib.rectangle(x_start, y_start, x_end, y_end)
        shape = shape_predictor(image, rect)

        face_descriptor = face_model.compute_face_descriptor(image, shape)
        feature_vector = np.array(face_descriptor)

        insert_query = "INSERT INTO feature_vectors (filename, vector) VALUES (%s, %s);"
        db_cursor.execute(insert_query, (filename, feature_vector.tolist()))
        db_conn.commit()

    db_cursor.close()
    db_conn.close()


if __name__ == '__main__':
    dataset_path = r"path_to_dataset_folder"
    extract_face_features(dataset_path)